import psutil
import tracemalloc
import threading
import time
import gc
import sys
from datetime import datetime
from typing import Dict, List, Optional, Callable
from functools import wraps
from sinotrans.utils import Logger
import json

class MemoryMonitor:
    """多维度内存监控器"""
    
    def __init__(self, 
                 memory_threshold_mb: int = 500,
                 warning_threshold_mb: int = 800,
                 critical_threshold_mb: int = 1000,
                 snapshot_interval: int = 30):
        
        self.memory_threshold_mb = memory_threshold_mb
        self.warning_threshold_mb = warning_threshold_mb
        self.critical_threshold_mb = critical_threshold_mb
        self.snapshot_interval = snapshot_interval
        
        # 监控数据存储
        self.memory_snapshots = []
        self.function_memory_stats = {}
        self.thread_memory_stats = {}
        self.variable_memory_stats = {}
        self.peak_memory_usage = 0
        self.memory_warnings = []
        
        # 控制标志
        self.monitoring_active = False
        self.snapshot_thread = None
        self._lock = threading.Lock()
        
        # tracemalloc相关
        self.tracemalloc_started = False
        self.baseline_snapshot = None
        
    def start_monitoring(self):
        """启动内存监控"""
        if not self.monitoring_active:
            self.monitoring_active = True
            
            # 启动tracemalloc
            if not tracemalloc.is_tracing():
                tracemalloc.start()
                self.tracemalloc_started = True
                self.baseline_snapshot = tracemalloc.take_snapshot()
                Logger.info("🔍 Tracemalloc监控已启动")
            
            # 启动定期快照线程
            self.snapshot_thread = threading.Thread(target=self._snapshot_worker, daemon=True)
            self.snapshot_thread.start()
            
            Logger.info(f"📊 内存监控已启动 - 阈值: {self.memory_threshold_mb}MB")
    
    def stop_monitoring(self):
        """停止内存监控"""
        self.monitoring_active = False
        
        if self.tracemalloc_started and tracemalloc.is_tracing():
            tracemalloc.stop()
            self.tracemalloc_started = False
            Logger.info("🔍 Tracemalloc监控已停止")
        
        if self.snapshot_thread:
            self.snapshot_thread.join(timeout=5)
        
        Logger.info("📊 内存监控已停止")
    
    def get_system_memory_info(self) -> Dict:
        """获取系统级内存信息(psutil)"""
        process = psutil.Process()
        memory_info = process.memory_info()
        memory_percent = process.memory_percent()
        
        # 系统内存信息
        system_memory = psutil.virtual_memory()
        
        return {
            'process_rss_mb': memory_info.rss / 1024 / 1024,  # 物理内存
            'process_vms_mb': memory_info.vms / 1024 / 1024,  # 虚拟内存
            'process_percent': memory_percent,
            'system_total_mb': system_memory.total / 1024 / 1024,
            'system_available_mb': system_memory.available / 1024 / 1024,
            'system_used_percent': system_memory.percent,
            'timestamp': datetime.now().isoformat()
        }
    
    def get_thread_memory_info(self) -> Dict:
        """获取线程级内存信息"""
        thread_info = {}
        current_thread = threading.current_thread()
        
        # 获取当前线程信息
        thread_info[current_thread.ident] = {
            'name': current_thread.name,
            'is_alive': current_thread.is_alive(),
            'daemon': current_thread.daemon
        }
        
        # 获取所有活跃线程
        for thread in threading.enumerate():
            if thread.ident not in thread_info:
                thread_info[thread.ident] = {
                    'name': thread.name,
                    'is_alive': thread.is_alive(),
                    'daemon': thread.daemon
                }
        
        return thread_info
    
    def take_tracemalloc_snapshot(self) -> Optional[Dict]:
        """获取tracemalloc快照"""
        if not tracemalloc.is_tracing():
            return None
        
        current_snapshot = tracemalloc.take_snapshot()
        
        # 与基线对比
        if self.baseline_snapshot:
            top_stats = current_snapshot.compare_to(self.baseline_snapshot, 'lineno')
        else:
            top_stats = current_snapshot.statistics('lineno')
        
        # 获取前10个内存分配热点
        hotspots = []
        for index, stat in enumerate(top_stats[:10]):
            hotspots.append({
                'rank': index + 1,
                'filename': stat.traceback.format()[-1] if stat.traceback else 'Unknown',
                'size_mb': stat.size / 1024 / 1024,
                'count': stat.count,
                'size_diff_mb': getattr(stat, 'size_diff', 0) / 1024 / 1024 if hasattr(stat, 'size_diff') else 0
            })
        
        return {
            'total_size_mb': sum(stat.size for stat in current_snapshot.statistics('filename')) / 1024 / 1024,
            'hotspots': hotspots,
            'timestamp': datetime.now().isoformat()
        }
    
    def check_memory_threshold(self) -> Dict:
        """检查内存阈值并记录警告"""
        memory_info = self.get_system_memory_info()
        current_memory_mb = memory_info['process_rss_mb']
        
        # 更新峰值内存
        if current_memory_mb > self.peak_memory_usage:
            self.peak_memory_usage = current_memory_mb
        
        warning_level = None
        if current_memory_mb > self.critical_threshold_mb:
            warning_level = 'CRITICAL'
        elif current_memory_mb > self.warning_threshold_mb:
            warning_level = 'WARNING'
        elif current_memory_mb > self.memory_threshold_mb:
            warning_level = 'INFO'
        
        if warning_level:
            warning = {
                'level': warning_level,
                'memory_mb': current_memory_mb,
                'threshold_mb': self.memory_threshold_mb,
                'timestamp': datetime.now().isoformat(),
                'thread_count': threading.active_count()
            }
            
            with self._lock:
                self.memory_warnings.append(warning)
                # 保持最近100条警告
                if len(self.memory_warnings) > 100:
                    self.memory_warnings = self.memory_warnings[-100:]
            
            Logger.warning(f"⚠️ 内存{warning_level}: {current_memory_mb:.1f}MB (阈值: {self.memory_threshold_mb}MB)")
        
        return {
            'current_memory_mb': current_memory_mb,
            'warning_level': warning_level,
            'peak_memory_mb': self.peak_memory_usage
        }
    
    def _snapshot_worker(self):
        """定期快照工作线程"""
        while self.monitoring_active:
            try:
                # 系统内存快照
                system_info = self.get_system_memory_info()
                thread_info = self.get_thread_memory_info()
                tracemalloc_info = self.take_tracemalloc_snapshot()
                
                snapshot = {
                    'system': system_info,
                    'threads': thread_info,
                    'tracemalloc': tracemalloc_info,
                    'timestamp': datetime.now().isoformat()
                }
                
                with self._lock:
                    self.memory_snapshots.append(snapshot)
                    # 保持最近100个快照
                    if len(self.memory_snapshots) > 100:
                        self.memory_snapshots = self.memory_snapshots[-100:]
                
                # 检查内存阈值
                self.check_memory_threshold()
                
            except Exception as e:
                Logger.error(f"❌ 内存快照失败: {e}")
            
            time.sleep(self.snapshot_interval)
    
    def force_garbage_collection(self) -> Dict:
        """强制垃圾回收"""
        before_memory = self.get_system_memory_info()['process_rss_mb']
        
        # 执行垃圾回收
        collected = gc.collect()
        
        after_memory = self.get_system_memory_info()['process_rss_mb']
        freed_mb = before_memory - after_memory
        
        Logger.info(f"🗑️ 垃圾回收完成: 释放 {freed_mb:.1f}MB, 回收对象 {collected}个")
        
        return {
            'before_memory_mb': before_memory,
            'after_memory_mb': after_memory,
            'freed_mb': freed_mb,
            'collected_objects': collected
        }
    
    def get_memory_report(self) -> Dict:
        """生成内存监控报告"""
        with self._lock:
            latest_snapshot = self.memory_snapshots[-1] if self.memory_snapshots else None
            recent_warnings = self.memory_warnings[-10:] if self.memory_warnings else []
        
        return {
            'monitoring_active': self.monitoring_active,
            'peak_memory_mb': self.peak_memory_usage,
            'current_snapshot': latest_snapshot,
            'recent_warnings': recent_warnings,
            'total_warnings': len(self.memory_warnings),
            'function_stats': dict(self.function_memory_stats),
            'variable_stats': dict(self.variable_memory_stats)
        }

# 函数级内存监控装饰器
def memory_profile(monitor_instance: MemoryMonitor, track_variables: List[str] = None):
    """函数级内存监控装饰器"""
    def decorator(func: Callable):
        @wraps(func)
        def wrapper(*args, **kwargs):
            if not monitor_instance.monitoring_active:
                return func(*args, **kwargs)
            
            func_name = f"{func.__module__}.{func.__qualname__}"
            thread_id = threading.current_thread().ident
            
            # 执行前内存状态
            before_memory = monitor_instance.get_system_memory_info()['process_rss_mb']
            before_time = time.time()
            
            try:
                # 执行函数
                result = func(*args, **kwargs)
                
                # 执行后内存状态
                after_memory = monitor_instance.get_system_memory_info()['process_rss_mb']
                after_time = time.time()
                
                # 记录函数内存统计
                memory_delta = after_memory - before_memory
                execution_time = after_time - before_time
                
                with monitor_instance._lock:
                    if func_name not in monitor_instance.function_memory_stats:
                        monitor_instance.function_memory_stats[func_name] = {
                            'call_count': 0,
                            'total_memory_delta_mb': 0,
                            'max_memory_delta_mb': 0,
                            'total_execution_time': 0,
                            'thread_calls': {}
                        }
                    
                    stats = monitor_instance.function_memory_stats[func_name]
                    stats['call_count'] += 1
                    stats['total_memory_delta_mb'] += memory_delta
                    stats['max_memory_delta_mb'] = max(stats['max_memory_delta_mb'], memory_delta)
                    stats['total_execution_time'] += execution_time
                    
                    # 线程级统计
                    if thread_id not in stats['thread_calls']:
                        stats['thread_calls'][thread_id] = 0
                    stats['thread_calls'][thread_id] += 1
                
                # 记录变量内存使用(如果指定)
                if track_variables and hasattr(func, '__code__'):
                    frame = sys._getframe(1)
                    for var_name in track_variables:
                        if var_name in frame.f_locals:
                            var_value = frame.f_locals[var_name]
                            var_size = sys.getsizeof(var_value)
                            
                            with monitor_instance._lock:
                                var_key = f"{func_name}.{var_name}"
                                if var_key not in monitor_instance.variable_memory_stats:
                                    monitor_instance.variable_memory_stats[var_key] = []
                                
                                monitor_instance.variable_memory_stats[var_key].append({
                                    'size_bytes': var_size,
                                    'timestamp': datetime.now().isoformat(),
                                    'thread_id': thread_id
                                })
                                
                                # 保持最近50条记录
                                if len(monitor_instance.variable_memory_stats[var_key]) > 50:
                                    monitor_instance.variable_memory_stats[var_key] = monitor_instance.variable_memory_stats[var_key][-50:]
                
                if memory_delta > 10:  # 内存增长超过10MB时记录
                    Logger.info(f"📈 函数 {func_name} 内存增长: {memory_delta:.1f}MB, 耗时: {execution_time:.2f}s")
                
                return result
                
            except Exception as e:
                Logger.error(f"❌ 函数 {func_name} 执行出错: {e}")
                raise
        
        return wrapper
    return decorator